package translationModel;

import java.util.ArrayList;
import java.util.BitSet;


public class PhraseDecoder 
{
	private Sentence f;
	private ArrayList<Phrase> phrases;
	private LanguageModel languageModel;
	private Sentence ref;
	private int beamSize;
	private int distortionLimit;
	private ArrayList<StateStack> stateStacks;
	
	public PhraseDecoder(Sentence f, ArrayList<Phrase> lexicon, LanguageModel lm, int beamSize, int distortionLimit)
	{
		this(f, lexicon, lm, null, beamSize, distortionLimit);
	}
	
	public PhraseDecoder(Sentence f, ArrayList<Phrase> lexicon, LanguageModel lm, Sentence ref)
	{
		this(f, lexicon, lm, ref, -1, -1);
	}
	
	private PhraseDecoder(Sentence f, ArrayList<Phrase> lexicon, LanguageModel lm, Sentence ref, int beamSize, int distortionLimit)
	{
		this.beamSize = beamSize;
		this.distortionLimit = distortionLimit;
		this.f = f;
		this.ref = ref;
		this.phrases = GetRelevantPhrases(lexicon);
		this.languageModel = lm;
		
		stateStacks = new ArrayList<StateStack>(f.length() + 1);
		for (int i = 0; i <= f.length(); i++)
			stateStacks.add(new StateStack(this.beamSize));
		
		State q0 = new State();
		q0.preLastWord = null;
		q0.lastWord = "<s>";
		q0.cover = new BitSet(f.length());
		q0.fEndPoint = -1;
		q0.eEndPoint = -1;
		q0.score = 0;
		stateStacks.get(0).AddState(q0);
	}
	
	private ArrayList<Phrase> GetRelevantPhrases(ArrayList<Phrase> lexicon)
	{
		ArrayList<Phrase> phrases = new ArrayList<Phrase>();
		for (Phrase ph : lexicon)
		{
			// Add the phrase only if it appears in the source, and, if needed, check also the reference
			if ((f.FindSubSentence(ph.f) != -1) && ((ref == null) || (ref.FindSubSentence(ph.e) != -1)))
				phrases.add(ph);
		}
		return phrases;
	}
	
	public String RunAlgorithm()
	{
		for (int i = 0; i < f.length(); i++)
		{
			stateStacks.get(i).Beam();
			for (State q : stateStacks.get(i).GetStates())
			{
				for (TranslatedPhrase p : PhrasesToAppend(q))
				{
					State next = NextState(q, p);
					stateStacks.get(next.cover.cardinality()).AddState(next);
				}
			}
		}
		
		State state = GetBestState();
		return GetTranslationString(state);
	}
	
	private ArrayList<TranslatedPhrase> PhrasesToAppend(State state)
	{
		ArrayList<TranslatedPhrase> phrasesToAppend = new ArrayList<TranslatedPhrase>();
		for (Phrase ph : phrases)
		{
			// Check if the phrase translation appears next in the reference, or there is no reference
			if ((ref == null) || (ref.FindSubSentence(state.eEndPoint + 1, ph.e) == state.eEndPoint + 1))
			{
				// For each appearance of phrase in the source
				for (int i = f.FindSubSentence(ph.f); i != -1; i = f.FindSubSentence(i + 1, ph.f))
				{
					// Check the distortion limit
					if ((distortionLimit == -1) || (Math.abs(state.fEndPoint + 1 - i) <= distortionLimit))
					{
						// Check if the phrase doesn't overlap the previous state
						boolean flagOverlap = false;
						for (int j = 0; j < ph.f.length(); j++)
						{
							if (state.cover.get(i + j))
							{
								flagOverlap = true;
								break;
							}
						}
						if (!flagOverlap)
							phrasesToAppend.add(new TranslatedPhrase(ph, i, i + ph.f.length() - 1, state.eEndPoint + 1, state.eEndPoint + ph.e.length()));
					}
				}
			}
		}
		return phrasesToAppend;
	}
	
	private State NextState(State q, TranslatedPhrase p)
	{
		State next = new State();
		if (p.phrase.e.length() > 1)
			next.preLastWord = p.phrase.e.getWord(p.phrase.e.length() - 2);
		else
			next.preLastWord = q.lastWord;
		next.lastWord = p.phrase.e.getWord(p.phrase.e.length() - 1);
		next.cover = new BitSet(f.length());
		next.cover.or(q.cover);
		for (int i = p.fStart; i <= p.fEnd; i++)
			next.cover.set(i);
		next.fEndPoint = p.fEnd;
		next.eEndPoint = p.eEnd;
		next.score = q.score + p.phrase.score;
		if (p.phrase.e.length() > 1)
			next.score += languageModel.GetScore(q.lastWord, p.phrase.e.getWord(0), p.phrase.e.getWord(1));
		if (q.preLastWord != null)
			next.score += languageModel.GetScore(q.preLastWord, q.lastWord, p.phrase.e.getWord(0));
		next.phrase = p;
		next.prevState = q;
		return next;
	}
	
	private State GetBestState()
	{
		for (int i = f.length(); i >= 0; i--)
		{
			StateStack current = stateStacks.get(i);
			if (!current.GetStates().isEmpty())
			{
				// Add the score of '</s>' from the language model
				AddSentenceEndScore(current);
				// If there is no reference, get the best score state
				if (ref == null)
					return current.GetBestScore();
				else
				{
					// There is a reference, so check if the state covers all the reference
					State state = current.GetBestEEnd(ref.length() - 1);
					if (state != null)
						return state;
				}
			}
		}
		return null;
	}
	
	private void AddSentenceEndScore(StateStack stack)
	{
		for (State state : stack.GetStates())
			state.score += languageModel.GetScore(state.preLastWord, state.lastWord, "</s>");
	}
	
	private String GetTranslationString(State state)
	{
		if ((state == null) || (state.cover.cardinality() == 0))
			return "Couldn't Find Translation!";
		// Backtrack the translation
		double score = state.score;
		StringBuilder translation = new StringBuilder();
		StringBuilder parts = new StringBuilder();
		while (state.cover.cardinality() > 0)
		{
			translation.insert(0, " " + state.phrase.phrase.e.toString());
			parts.insert(0, " (" + state.phrase.eStart + "," + state.phrase.eEnd + ")");
			state = state.prevState;
		}
		translation.deleteCharAt(0);
		parts.deleteCharAt(0);
		return translation.toString() + " ||| " + parts.toString() + " ||| " + score;
	}
} 
